import numpy as np

def distance(a, b):
    """计算两个向量之间的欧式距离"""
    squared_sum = 0.0
    for i in range(len(a)):
        squared_sum += (a[i] - b[i]) ** 2
    return np.sqrt(squared_sum)


class KNN:
    def __init__(self, k, label_num):
        """初始化 KNN 模型"""
        self.k = k
        self.label_num = label_num

    def fit(self, x_train, y_train):
        """在类中保存训练数据"""
        self.x_train = np.array(x_train)
        self.y_train = np.array(y_train)

    def get_knn_indices(self, x):
        """获取距离目标样本点最近的K个样本点的索引"""
        # 使用 distance 函数计算所有训练样本到目标样本的距离
        distances = [distance(train_sample, x) for train_sample in self.x_train]
        distances = np.array(distances)
        
        # 按距离从小到大排序，并得到对应的下标
        knn_indices = np.argsort(distances)
        # 取最近的K个
        knn_indices = knn_indices[:self.k]
        return knn_indices

    def get_label(self, x):
        """观察K个近邻并获取其中数量最多的类别"""
        knn_indices = self.get_knn_indices(x)
        # 类别计数
        label_statistic = np.zeros(self.label_num)
        for index in knn_indices:
            label = int(self.y_train[index])
            label_statistic[label] += 1
        # 返回数量最多的类别
        return np.argmax(label_statistic)

    def predict(self, x_test):
        """预测样本 x_test 的类别"""
        x_test = np.array(x_test)
        predicted_test_labels = []
        for i, x in enumerate(x_test):
            predicted_test_labels.append(self.get_label(x))
        return np.array(predicted_test_labels)


if __name__ == "__main__":
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split

    # 加载数据集
    digits = load_digits()
    x_data = digits.data
    y_data = digits.target

    # 划分训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(
        x_data, y_data, test_size=0.3, random_state=42)

    # 测试不同的 k 值
    print("测试不同 K 值的预测准确率：")
    for k in [1, 3, 5]:
        knn = KNN(k, label_num=10)
        knn.fit(x_train, y_train)
        predicted_labels = knn.predict(x_test)

        accuracy = np.mean(predicted_labels == y_test)
        print(f'K的取值为 {k}, 预测准确率为 {accuracy * 100:.1f}%')
